#!/usr/bin/perl -w

use strict;
use Getopt::Long;

my $plinkfile = '';      # plink file name, not including extension
my $ifBinary = '';       # 1 if the file set is binary, 0 if not
my $ifSexCoded = '';     # 1 if the gender is coded (1male,2female), 0 if not coded
my $markerCallRate = ''; # cutoff for marker call rate, eg. 0.03
my $sampleCallRate = ''; # cutoff for sample call rate, eg. 0.03
my $mafCutoff = '';      # cutoff for minor allele freq, eg. 0.01
my $hweCutoff = '';      # cutoff for hwe, eg. 0.00001
GetOptions('file=s' => \$plinkfile,
		   'b' => \$ifBinary,
		   's' => \$ifSexCoded,
		   'marker|m=s' => \$markerCallRate,
		   'sample|i=s' => \$sampleCallRate,
		   'maf|f=s' => \$mafCutoff,
		   'hwe|h=s' => \$hweCutoff);

####
# need to validate argument
####
		   
# print "$plinkfile\n$ifBinary\n$ifSexCoded\n$markerCallRate\n$sampleCallRate\n$hweCutoff\n";

# ##########################################################		
# # 1.if the file set is not binary, call plink --make-bed
# ##########################################################
# if (!$ifBinary){
	# system("plink --file $plinkfile --make-bed --out $plinkfile");
# }

# ##########################################################
# # 2.log sample number, marker count in log
# ##########################################################		   
# system("plink --bfile $plinkfile --out $plinkfile.summary");
# open LOG,"$plinkfile.summary.log";
# open REPORT,">$plinkfile.QCreport";
# while (<LOG>){
	# print REPORT "there is a total of $1 markers\n" if /^(\d+)\smarkers/; # log number of markers
	# print REPORT "there is a total of $1 samples\n" if /^(\d+)\sindividuals\sread/;
# }
# print REPORT "\n";
# close LOG;
# close REPORT;
# ##########################################################
# # 3. log number of markers in each chromosome (1-25)
# ##########################################################
# open REPORT,">>$plinkfile.QCreport";
# for (my $i = 1; $i != 26; $i++){
	# system("plink --bfile $plinkfile --chr $i --out $plinkfile.chr$i");
	# open LOG,"$plinkfile.chr$i.log";
	# while (<LOG>){
		# print REPORT "$i\t$1\n" if /^After\sfrequency\sand\sgenotyping\spruning,\sthere\sare\s(\d+)\sSNPs/;
	# }
	# close LOG;
# }
##########################################################
# 3. Sex check
##########################################################
# system("plink --bfile $plinkfile --impute-sex --make-bed --out $plinkfile.sexcoded");

# my @sexfailed_fid;
# my @sexfailed_id; # array of sample IDs that failed the sex check
# my @sexfailed_reported; # reported sex 
# my @sexfailed_imputed; # imputed sex
# my @sexfailed_ambigious; # array of sample IDs that have ambigious sex

# my @badsamples; #bad samples that need to be removed (imputed sex = 0)
# open SEXCHECK,"$plinkfile.sexcoded.sexcheck";
# while (<SEXCHECK>){
	# if (/PROBLEM/){
		# push(@sexfailed_id,((split)[0]."\t".(split)[1]));
		# push(@sexfailed_reported,(split)[2]);
		# push(@sexfailed_imputed,(split)[3]);
		# push(@sexfailed_ambigious,((split)[0]."\t".(split)[1])) if ((split)[3] == 0); # if imputed sex code is 0, sample sex is ambigious
	# }
# }
# close SEXCHECK;
# open REPORT,">>$plinkfile.QCreport";
# if (@sexfailed_id == 0){
	# print REPORT "no sample failed sex check\n\n";
# }
# else {
	# print REPORT "IID\tREPORTED\tIMPUTED\n"; # print header line
	# if ($ifSexCoded){
		# for(my $i = 0; $i < scalar @sexfailed_id ; $i++){
			# print REPORT "$sexfailed_id[$i]\t$sexfailed_reported[$i]\t$sexfailed_imputed[$i]\n" if (($sexfailed_reported[$i] != $sexfailed_imputed[$i])||$sexfailed_imputed[$i] == 0);
		# }
	# }
	# else {
		# for(my $i = 0; $i < scalar @sexfailed_id ; $i++){
			# print REPORT "$sexfailed_id[$i]\t$sexfailed_reported[$i]\t$sexfailed_imputed[$i]\n";
		# }
	# }
	# print REPORT "\n";
	# # print out the ambigious samples
	# for (my $i = 0; $i < scalar @sexfailed_ambigious ; $i++){
		# print REPORT "$sexfailed_ambigious[$i] has ambigious sex\n";
	# }
# }

# # Samples with ambigious sex will be written to a file
# open AMBIGIOUS_SEX,">ambigious_sex.txt";
# foreach (@sexfailed_ambigious){
	# print AMBIGIOUS_SEX "$_\n";
	# print REPORT "$_ is removed from the analysis due to ambigious sex code\n";
# }

# close REPORT;
# # remove ambigious sex samples
# system("plink --bfile $plinkfile.sexcoded --remove ambigious_sex.txt --make-bed --out $plinkfile.failedsex_removed"); 

###########################################################
# 4. Call rate filtering
###########################################################

# Y chromosome is removed and all its markers are kept in an separate file set .Y
system("plink --bfile $plinkfile.failedsex_removed --chr 25 --write-snplist --make-bed --out $plinkfile.Y");
system("plink --bfile $plinkfile.failedsex_removed --exclude $plinkfile.Y.snplist --make-bed --out $plinkfile.noY");

# perform missingness to detect markers that fails completely.
system("plink --bfile $plinkfile.noY --missing --out $plinkfile.noY.missing");

# parse the marker missingness file .lmiss to look for markers that failed completely.
# marker fails completely if F_MISS is 1
open LMISS,"$plinkfile.noY.missing.lmiss";
open TEMP,">$plinkfile.noY.missing.snplist";
while (<LMISS>){
	next if /F_MISS/;
	print TEMP (split)[1]."\n" if ((split)[4] == 1); # print marker ID if its missing frequency is 1
}
close TEMP;
# remove markers from .missing.snplist from .noY dataset
system("plink --bfile $plinkfile.noY --exclude $plinkfile.noY.missing.snplist --make-bed --out $plinkfile.noY.nomissing");

# filter sample call rate and marker call rate based on specified threshold
system("plink --bfile $plinkfile.noY.nomissing --mind $sampleCallRate --make-bed --out $plinkfile.noY.sampleFilter");
system("plink --bfile $plinkfile.noY.sampleFilter --geno $markerCallRate --make-bed --out $plinkfile.noY.filtered");

# repeat the same procedure for Y only data set
# perform missingness to detect markers that fails completely.
system("plink --bfile $plinkfile.Y --missing --out $plinkfile.Y.missing");

# parse the marker missingness file .lmiss to look for markers that failed completely.
# marker fails completely if F_MISS is 1
open LMISS,"$plinkfile.Y.missing.lmiss";
open TEMP,">$plinkfile.Y.missing.snplist";
while (<LMISS>){
	next if /F_MISS/;
	print TEMP (split)[1]."\n" if ((split)[4] == 1); # print marker ID if its missing frequency is 1
}
close TEMP;
# remove markers from .missing.snplist from .Y dataset
system("plink --bfile $plinkfile.Y --exclude $plinkfile.Y.missing.snplist --make-bed --out $plinkfile.Y.nomissing");

# filter sample call rate and marker call rate based on specified threshold
system("plink --bfile $plinkfile.Y.nomissing --mind $sampleCallRate --make-bed --out $plinkfile.Y.sampleFilter");
system("plink --bfile $plinkfile.Y.sampleFilter --geno $markerCallRate --make-bed --out $plinkfile.Y.filtered");

# merge noY and Y data set
system("plink --bfile $plinkfile.noY.filtered --bmerge $plinkfile.Y.filtered.bed $plinkfile.Y.filtered.bim $plinkfile.Y.filtered.fam --make-bed --out $plinkfile.callrate.filtered");

# writing report on call rate filtering
open REPORT,">>$plinkfile.QCreport";
# number of marker on Y
open TEMP,"$plinkfile.Y.log";
while (<TEMP>){
	print REPORT "there are $1 markers on Chr Y\n" if /^After frequency and genotyping pruning, there are (\d+) SNPs\n/; 
	print REPORT "the average call rate for the Y chr SNPs was $1\n" if /^Total genotyping rate in remaining individuals is (\d+\.\d+)/;
}
close TEMP;
# number of marker on 1-22,X,XY
open TEMP,"$plinkfile.noY.log";
while (<TEMP>){
	print REPORT "there are $1 markers on non-Chr Y\n" if /^Before frequency and genotyping pruning, there are (\d+) SNPs\n/; 
	print REPORT "the average call rate for the non-Y SNPs was $1\n" if /^Total genotyping rate in remaining individuals is (\d+\.\d+)/;
}	
close TEMP;
# number of failed marker on non-Y chr
open TEMP,"$plinkfile.noY.missing.snplist";
my $noYfailedsnp = 0;
while (<TEMP>){
	$noYfailedsnp++;
}
close TEMP;
print REPORT "$noYfailedsnp SNPs completely failed in non-Y chromosomes\n";
# report failed sample
open TEMP,"$plinkfile.noY.sampleFilter.irem";
while (<TEMP>){
	chomp;
	print REPORT "$_ failed $sampleCallRate sample call rate filter\n";
}
close TEMP;
# report total avg call rate, sample number and marker count on non-Y
open TEMP,"$plinkfile.noY.filtered.log";
while (<TEMP>){
	print REPORT "the average call rate for non-Y chromosome SNPs after $markerCallRate call rate filtering was $1\n" if /^Total genotyping rate in remaining individuals is (\d+\.\d+)/;
	print REPORT "$1 markers remain after call rate filtering in non Y chr\n" if /^After frequency and genotyping pruning, there are (\d+) SNPs\n/;
	print REPORT ($1+$2+$3)." samples remain after call rate filtering in non Y chr\n" if /After filtering, (\d+) cases, (\d+) controls and (\d+) missing/;
}
close TEMP;

#report on Y chr

# number of failed marker on Y chr
open TEMP,"$plinkfile.Y.missing.snplist";
my $Yfailedsnp = 0;
while (<TEMP>){
	$Yfailedsnp++;
}
close TEMP;
print REPORT "$Yfailedsnp SNPs completely failed in Y chromosomes\n";
# report total avg call rate, marker count on Y
open TEMP,"$plinkfile.Y.filtered.log";
while (<TEMP>){
	print REPORT "the average call rate for Y chromosome SNPs after $markerCallRate call rate filtering was $1\n" if /^Total genotyping rate in remaining individuals is (\d+\.\d+)/;
	print REPORT "$1 markers remain after call rate filtering in Y chr\n" if /^After frequency and genotyping pruning, there are (\d+) SNPs\n/;
}
close TEMP;


# report total avg call rate, marker count on all chromosomes
open TEMP,"$plinkfile.callrate.filtered.log";
while (<TEMP>){
	print REPORT "\nOverall, average call rate for SNPs after $markerCallRate call rate filtering was $1\n" if /^Total genotyping rate in remaining individuals is (\d+\.\d+)/;
	print REPORT "$1 markers remain after $sampleCallRate call rate filtering in Y chr\n" if /^After frequency and genotyping pruning, there are (\d+) SNPs\n/;
}
close TEMP;

##########################################################
# 5. Minor Allele Frequency (maf) filtering
##########################################################

system("plink --bfile $plinkfile.callrate.filtered --maf $mafCutoff --make-bed --out $plinkfile.maf.filtered");

open TEMP,"$plinkfile.maf.filtered.log";
while (<TEMP>){
	print REPORT "$1 markers remain after minor allele frequency filtering\n" if /^After frequency and genotyping pruning, there are (\d+) SNPs\n/;
	print REPORT "\n$1 markers had observed minor allele frequency lower than $mafCutoff and were removed from subsequent analysis\n" if /(\d+) SNPs failed frequency test/;
}
close TEMP;

##########################################################
# 6. Hardy-Weinberg equilibrium (option needs to be added 
# to do filtering on control only or all samples
##########################################################
system ("plink --bfile $plinkfile.maf.filtered --hwe $hweCutoff --make-bed --out $plinkfile.hwe.filtered");

open TEMP,"$plinkfile.hwe.filtered.log";
while (<TEMP>){
	print REPORT "$1 markers remain after hwe filtering\n" if /^After frequency and genotyping pruning, there are (\d+) SNPs\n/;
	print REPORT "\n$1 markers were out of Hardy-Weinberg equilibirium \($hweCutoff\) and were removed from subsequent analysis\n" if /(\d+) markers to be excluded based on HWE test/;
}
close TEMP;

